#! /usr/bin/env python3
import pandas
import pyspark.sql
from abc import abstractmethod
from enum import Enum


class DatasetType(Enum):
    SPARK_DATA_FRAME = 0
    PANDAS_DATA_FRAME = 1


class Dataset(object):
    """A ``Dataset`` represents a collection of structured data. A ``Dataset`` may
    be a wrapper of data object from other platforms such as Pandas and Spark.
    In this case, we can extract the underlying data object with method ``unwrap()``.

    >>> from gai.v2.unify import Dataset
    >>> import pandas as pd
    >>> df = pd.DataFrame({'col1': [1, 3], 'col2': [2, 4]})
    >>> ds = PandasDataFrame(df)
    >>> ds.get_type()
    <DatasetType.PANDAS_DATA_FRAME: 1>
    """

    @abstractmethod
    def get_type(self) -> DatasetType:
        """

        Returns:
			 the type indicator (``DatasetType``) corresponding to the underlying variant.
        """
        raise NotImplementedError()

    @abstractmethod
    def unwrap(self):
        """

        Returns:
			 the underlying data object
        """
        raise NotImplementedError()

    @abstractmethod
    def count(self):
        """

        Returns:
            the size of the dataset
        """
        raise NotImplementedError()


class SparkDataFrame(Dataset):
    """A wrapper of ``pyspark.sql.DataFrame``.
    """

    def __init__(self, df: pyspark.sql.DataFrame):
        super(SparkDataFrame, self).__init__()
        assert isinstance(df, pyspark.sql.DataFrame)
        self._df = df

    def get_type(self) -> DatasetType:
        return DatasetType.SPARK_DATA_FRAME

    def unwrap(self):
        return self._df

    def count(self):
        return self._df.count()


class PandasDataFrame(Dataset):
    """A wrapper of ``pandas.DataFrame``.
    """

    def __init__(self, df: pandas.DataFrame):
        super(PandasDataFrame, self).__init__()
        assert isinstance(df, pandas.DataFrame)
        self._df = df

    def get_type(self) -> DatasetType:
        return DatasetType.PANDAS_DATA_FRAME

    def unwrap(self):
        return self._df

    def count(self):
        return self._df.shape[0]